package com.lan.ai.controller;

import cn.hutool.core.util.ObjectUtil;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.lan.ai.mutimodelEmbedding.MutimodeDocument;
import io.swagger.v3.oas.annotations.tags.Tag;
import lombok.RequiredArgsConstructor;
import lombok.SneakyThrows;
import org.springframework.ai.document.Document;
import org.springframework.ai.model.Media;
import org.springframework.ai.vectorstore.MilvusVectorStore;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.core.io.InputStreamResource;
import org.springframework.http.MediaType;
import org.springframework.util.MimeTypeUtils;
import org.springframework.web.bind.annotation.*;

import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author lanyanhua
 * @date 2024/12/24 15:57
 * @description
 */
@RestController
@RequestMapping("/ai/vector")
@Tag(name = "多模态向量数据库")
@RequiredArgsConstructor
public class MutimodeVectorStoreController {

    private final MilvusVectorStore imgVectorStore;

    @SneakyThrows
    @PostMapping("addImg")
    public Boolean add(@RequestParam(required = false) String content, @RequestParam(required = false) String file) {

        List<Media> objects = ObjectUtil.isNotEmpty(file) ?
                List.of(new Media(new MediaType(MimeTypeUtils.IMAGE_PNG), new InputStreamResource(new FileInputStream(file))))
                : List.of();
        List<Document> documentList = List.of(new MutimodeDocument(content == null ? "" : content,
                objects,
                Map.of("file",file))

        );
        imgVectorStore.add(documentList);
        return true;
    }

    /**
     * 查询向量数据库
     *
     * @return 匹配到的文档
     */
    @GetMapping("imageFilter")
    public List<Document> imageFilter(String content, String file) throws FileNotFoundException {

        List<Media> objects = ObjectUtil.isNotEmpty(file) ?
                List.of(new Media(new MediaType(MimeTypeUtils.IMAGE_PNG), new InputStreamResource(new FileInputStream(file))))
                : List.of();
        MutimodeDocument document = new MutimodeDocument(content == null ? "" : content,
                objects,
                Map.of());

        return imgVectorStore.similaritySearch(SearchRequest.query(document.getContent()).withSimilarityThreshold(0.3d).withTopK(6));
    }

}
